'''
docker run --rm -p 12455:12455 -v /home/xs/AI_docker/app:/home/app -v /home/xs/AI_docker/models:/home/models --name tensorrt  -it --runtime=nvidia baoxin/tensorrt /bin/bash start.sh 1
'''
from fastapi import FastAPI
from pydantic import BaseModel
import uvicorn
import cv2
import random
import os
import base64
import json
import time
import numpy as np
from enum import Enum

#加载前面的检测类
from project_deploy import *
from model_ways import *
from post_media import *
from model_utils import *


app = FastAPI()

def base64toCv(base64_src):
    img_b64decode = base64.b64decode(base64_src)  # base64解码
    img_array = np.fromstring(img_b64decode, np.uint8)  # 转换np序列
    img_cv = cv2.imdecode(img_array, cv2.COLOR_BGR2RGB)  # 转换OpenCV格式
    return img_cv

# 将识别的类别加入枚举
class Targets(str, Enum):
    det_target1 = "cap"

# 定义接收数据的结构
class Item(BaseModel):
    base64: str = None  #图片base64
    target: Targets = None   #识别类型


@app.post('/detector')
async def calculate(request_data: Item):
    Target = request_data.target
    img_base64 = request_data.base64
    cam_num=1
    # 输出检测的信息和调用时间
    print("Detection for", Target.value , "! Time:", time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))


   # frame = base64toCv(img_base64)
    frame = cv2.imread('picture/test1.jpg')
    frame = cv2.resize(frame, (608, 608))
    locals()['frame'+str(cam_num-1)]=electron_frame(frame,(608, 608),electron_fences[cam_num-1]) 

    #frame初始化,yolov4所需预处理
    IN_IMAGE_H, IN_IMAGE_W = (608,608)
    resized = cv2.resize(frame, (IN_IMAGE_W, IN_IMAGE_H), interpolation=cv2.INTER_LINEAR)
    img_in = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB)
    img_in = np.transpose(img_in, (2, 0, 1)).astype(np.float32)
    img_in = np.expand_dims(img_in, axis=0)
    img_in /= 255.0
    img_in = np.ascontiguousarray(img_in)
    if cam_num==1:
        img_in_all=img_in
    model.detect_frame_all(img_in_all)

    frame,result=model.detect_frame_pic(locals()['frame'+str(cam_num-1)],cam_num)
    results= {}
    for i in range(0,len(result)):
        res={i:[result[i][0],result[i][1],int(result[i][2]*100),result[i][3]]}
        results.update(res)
    return results

   # return {"code": 201, "data": []}


if __name__ == '__main__':
    # 实例化一个检测器
    model_name='cap'
    gpu_id=0
    id_list=['1']
    electron_fences=['None']
    ffmpeg_types=['picture']

    model=ModelFun(model_name,gpu_id,id_list,electron_fences,ffmpeg_types)
    image_size = (eval(model_name+'.size'), eval(model_name+'.size'))
    print('Loaded yolo model!')
    print("Service start!")
    uvicorn.run(app=app,
        host="0.0.0.0", # 服务器填写0.0.0.0
        port=12455,
        workers=1)


